/** @file

  Copyright (C) 2022 - 2023, Phytium Technology Co., Ltd. All rights reserved.<BR>

  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <Library/PhytiumSpiNorFlashLib.h>
#include <Library/DebugLib.h>
#include <Protocol/SpiFlash.h>
#include <Library/IoLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>

/**
  Erase the flash.

  @param[in]  This    Pointer of  EFI_SPI_FLASH_PROTOCOL.
  @param[in]  Offset    The start address to erase
  @param[in]  Length    The length to erase

  @retval EFI_SUCCESS   Function successful returned.
  @retval EFI_INVALID_PARAMETER   The input parameters is not valid.
**/
EFI_STATUS
Erase (
  IN EFI_SPI_FLASH_PROTOCOL  *This,
  IN UINT64                  Offset,
  IN UINT64                  Length
  )
{
  EFI_STATUS     Status;
  UINT64         Index;
  UINT64         Count;

  Status = EFI_SUCCESS;
  if ((Length % SIZE_64KB) == 0) {
    Count = Length / SIZE_64KB;
    for (Index = 0; Index < Count; Index ++) {
      NorFlashPlatformEraseSingleBlock (Offset);
      Offset += SIZE_64KB;
    }
  } else {
    Status = EFI_INVALID_PARAMETER;
  }

  return Status;
}

/**
  Write data buffer to the flash.

  @param[in]  This      Pointer of  EFI_SPI_FLASH_PROTOCOL.
  @param[in]  Offset    The start address to write.
  @param[in]  Buffer    The data buffer to write.
  @param[in]  Length    The length to write.

  @retval EFI_SUCCESS   Function successful returned.
  @retval EFI_INVALID_PARAMETER   The input parameters is not valid.
**/
EFI_STATUS
Write (
  IN EFI_SPI_FLASH_PROTOCOL  *This,
  IN  UINT64                 Offset,
  IN  UINT8                  *Buffer,
  IN  UINT64                 Length
  )
{
  EFI_STATUS     Status;

  if (((Length % 4) != 0)|| ((Offset % 4) != 0)) {
    DEBUG ((EFI_D_ERROR, "Error input for spi program!!Need check!!\n"));
    return EFI_INVALID_PARAMETER;
  }
  Status = NorFlashPlatformWrite(Offset, Buffer, Length);

  return Status;
}

/**
  Read the flash data to buffer.

  @param[in]  This      Pointer of  EFI_SPI_FLASH_PROTOCOL.
  @param[in]  Offset    The start address to read.
  @param[in,out]  Buffer    The data buffer to read.
  @param[in]  Length    The length to read.

  @retval EFI_SUCCESS   Function successful returned.
**/
EFI_STATUS
Read (
  IN EFI_SPI_FLASH_PROTOCOL  *This,
  IN UINT64                  Offset,
  IN OUT UINT8               *Buffer,
  IN UINT64                  Length
  )
{

  NorFlashPlatformRead (Offset, Buffer, Length);

  return EFI_SUCCESS;
}

/**
  Erase the flash than write values.

  @param[in]  This      Pointer of  EFI_SPI_FLASH_PROTOCOL.
  @param[in]  Offset    The start address to erase and write.
  @param[in]  Buffer    The data buffer to write.
  @param[in]  Length    The length to erase and write.

  @retval EFI_SUCCESS   Function successful returned.
  @retval EFI_INVALID_PARAMETER   The input parameters is not valid.
**/
EFI_STATUS
EraseWrite (
  IN EFI_SPI_FLASH_PROTOCOL  *This,
  IN  UINT64                 Offset,
  IN  UINT8                  *Buffer,
  IN  UINT64                 Length
  )
{
  EFI_STATUS      Status;
  UINTN           Index;
  UINTN           Count;

  Status = EFI_SUCCESS;
  Count = Length / SIZE_64KB;
  for (Index = 0; Index < Count; Index ++) {
    Status = Erase (This, Offset + Index * SIZE_64KB, SIZE_64KB);
    if (EFI_ERROR(Status)) {
      break;
    }
    Status = Write (This, Offset + Index * SIZE_64KB, Buffer + Index * SIZE_64KB, SIZE_64KB);
    if (EFI_ERROR(Status)) {
      break;
    }
  }
  return Status;
}

/**
  Write data buffer to the flash.

  @param[in]  This      Pointer of  EFI_SPI_FLASH_PROTOCOL.
  @param[in]  Offset    The start address to write.
  @param[in]  Buffer    The data buffer to write.
  @param[in]  Size      The length to write.

  @retval EFI_SUCCESS   Function successful returned.
  @retval EFI_INVALID_PARAMETER   The input parameters is not valid.
**/
EFI_STATUS
KunlunFlashWrite(
  IN EFI_SPI_FLASH_PROTOCOL  *This,
  IN  UINT64                 Offset,
  IN  UINT8                  *Buffer,
  IN  UINT64                 Size
  )
{
  EFI_STATUS              Status;
  UINTN                   BlockSize;
  UINT8                   *FlashDataTemp;
  UINT32                  BlockAddrStart;
  UINT16                  CurOffset;
  UINT16                  BlockLen;
  UINT64                  Length;

  Status  = EFI_SUCCESS;
  if (Size == 0) {
    return EFI_INVALID_PARAMETER;
  }
  BlockSize  = 0x010000;
  FlashDataTemp = AllocateZeroPool(BlockSize);
  if (FlashDataTemp == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  Length = Size;
  CurOffset = Offset % BlockSize;
  BlockAddrStart = (Offset / BlockSize) * BlockSize;
  while(Length > 0)
  {
    CopyMem(FlashDataTemp,(UINT8 *)((UINTN)BlockAddrStart),BlockSize);

    Status = Erase(This,BlockAddrStart, BlockSize);
    if (Status != EFI_SUCCESS) {
      DEBUG((DEBUG_ERROR, "%a(%d) Flash Erase Error.\n", __FILE__, __LINE__));
      return Status;
    }

    BlockLen = (CurOffset + Length) > BlockSize ? (BlockSize - CurOffset) : Length;
    CopyMem(FlashDataTemp+CurOffset, Buffer, BlockLen);

    Length -= BlockLen;
    Buffer = Buffer + BlockLen;
    Status = EraseWrite(This,BlockAddrStart, FlashDataTemp, BlockSize);
    if (Status != EFI_SUCCESS) {
      DEBUG((DEBUG_ERROR, "%a(%d) Flash EraseWrite Error.\n", __FILE__, __LINE__));
      return Status;
    }

    BlockAddrStart += BlockSize;
    CurOffset = 0;
  }

  return Status;
}

EFI_SPI_FLASH_PROTOCOL  gSpiFlashProtocol = {
  Erase,
  KunlunFlashWrite,
  Read,
  EraseWrite
};
